#!/usr/bin/env python

from argparse import ArgumentParser
import collections
import cv2
from dt_apriltags import Detector
def _camera_params(pstr):

    print(pstr)
    pstr = pstr.strip()

    if pstr[0] == '(' and pstr[-1] == ')':
        pstr = pstr[1:-1]

    params = tuple( [ float(param.strip()) for param in pstr.split(',') ] )

    assert( len(params) ==  4)

    return params


def main():

    '''Main function.'''

    parser = ArgumentParser(
        description='apriltag detection')

    parser.add_argument('device_or_movie', metavar='INPUT', nargs='?', default=0,
                        help='Movie to load or integer ID of camera device')

    parser.add_argument('-k', '--camera-params', type=_camera_params,
                        default=None,
                        help='intrinsic parameters for camera (in the form fx,fy,cx,cy)')

    parser.add_argument('-s', '--tag-size', type=float,
                        default=1.0,
                        help='tag size in user-specified units (default=1.0)')

    parser.add_argument('-v', '--visible', type=bool,
                        default=False,
                        help='show camera image or not')



    options = parser.parse_args()

    print(options)

    try:
        cap = cv2.VideoCapture(int(options.device_or_movie))
    except ValueError:
        cap = cv2.VideoCapture(options.device_or_movie)

    while True:
        t1 = cv2.getTickCount()
        success, frame = cap.read()
        if not success:
            print('cam read failed')
            continue

        tags = april_detect(frame, options)

        t2 = cv2.getTickCount()
        time1 = (t2-t1)/freq
        frame_rate_calc= 1/time1

        if options.visible:
            cv2.putText(frame,'FPS: {0:.2f}'.format(frame_rate_calc),(30,50),cv2.FONT_HERSHEY_SIMPLEX,1,(255,255,0),2,cv2.LINE_AA)

            april_show_detection(frame,tags)

            cv2.imshow("apriltag",frame)

        print('FPS: {0:.2f}'.format(frame_rate_calc))

        k = cv2.waitKey(1)

        if k == ord('q'):
            break


# Initialize frame rate calculation
frame_rate_calc = 1
freq = cv2.getTickFrequency()

at_detector = Detector(families='tag36h11',
                   nthreads=4,
                   quad_decimate=1.0,
                   quad_sigma=0.0,
                   refine_edges=1,
                   decode_sharpening=0.25,
                   debug=0)


def april_detect(frame,options):

    gray = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY)

    tags = at_detector.detect(gray, True, options.camera_params, options.tag_size)

    return tags

def april_show_detection(frame, tags):
    for tag in tags:
        for idx in range(len(tag.corners)):
            cv2.line(frame, tuple(tag.corners[idx-1, :].astype(int)),
                    tuple(tag.corners[idx, :].astype(int)), (0, 255, 0))

        cv2.putText(frame,
                "[{}]\n{}".format(str(tag.tag_id),
                    tag.pose_t),
                org=(tag.corners[0, 0].astype(int)-100,tag.corners[0, 1].astype(int)+10),
                fontFace=cv2.FONT_HERSHEY_SIMPLEX,
                fontScale=0.5,
                color=(0, 0, 255))

if __name__ == '__main__':
    main()
